# 堆叠去噪自动编码机模型类
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class MpeSdaeModel(nn.Module):
    def __init__(self, input_dim, hidden_dim1, hidden_dim2, latent_dim):
        super(MpeSdaeModel, self).__init__()
        # 编码器部分
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim1),
            nn.ReLU(),
            nn.Linear(hidden_dim1, hidden_dim2),
            nn.ReLU(),
            nn.Linear(hidden_dim2, latent_dim)
        )
        
        # 解码器部分
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim2),
            nn.ReLU(),
            nn.Linear(hidden_dim2, hidden_dim1),
            nn.ReLU(),
            nn.Linear(hidden_dim1, input_dim),
            nn.Sigmoid()  # 通常用于输出层，因为输出值需要在[0, 1]之间
        )

    def forward(self, x):
        # 编码
        encoded = self.encoder(x)
        # 解码
        decoded = self.decoder(encoded)
        return decoded

    @staticmethod
    def add_noise(inputs, noise_factor=0.3):
        noise = torch.randn_like(inputs) * noise_factor
        noisy_inputs = inputs + noise
        return noisy_inputs.clamp(0.0, 1.0)  # 确保噪声数据在[0, 1]范围内

